import argparse
import json
import os
import sys
import ipe_param_face_recongnition

sys.path.append("../../../")

from common.python.common import *
from common.python.callback_data_struct import *
from common.python.infer_process import *
from common.python import image_decoder


def write_callback(args):
    p_host_infer_buf_out, f, output_size = args
    np_out = sdk.lyn_ptr_to_numpy(
        p_host_infer_buf_out, (output_size,), sdk.lyn_data_type_t.DT_UINT8
    )
    f.write(np_out.tobytes())
    sdk.c_free(p_host_infer_buf_out)
    return 0


def main():
    parser = argparse.ArgumentParser(description="Face Feature Extraction")
    parser.add_argument("face_image_path", type=str, help="input face image path")
    parser.add_argument("config_file", type=str, help="input face config file")
    parser.add_argument("model_path", type=str, help="face feature model path")
    parser.add_argument("device_id", type=int, help="lynxi device id")

    args = parser.parse_args()

    face_image_path = args.face_image_path
    config_file = args.config_file
    model_path = args.model_path
    device_id = args.device_id

    if not os.path.exists(face_image_path):
        print("error: input face image path does not exist.")
        return

    if not os.path.exists(config_file):
        print("error: input face config file does not exist.")
        return

    if not os.path.exists(model_path):
        print("error: face feature model path does not exist.")
        return

    if not is_valid_device(device_id):
        print("error: device id is invalid.")
        return

    with open(config_file) as f:
        js = json.load(f)

    feature_path = js["features"]
    images = js["images"]

    # 2. 创建 context 与 stream
    ctx, ret = sdk.lyn_create_context(device_id)
    error_check(ret, "create context")
    ret = sdk.lyn_register_error_handler(default_stream_error_handle)
    error_check(ret, "lyn_register_error_handler")
    stream, ret = sdk.lyn_create_stream()
    error_check(ret, "lyn_create_stream")

    # 3. 初始化
    model_info = ModelInfo(model_path)

    ipe = ipe_param_face_recongnition.IpeParamW600k_mbf(
        model_info.width, model_info.height
    )

    p_decode_buf_out, ret = sdk.lyn_malloc(int(round(1920 * 1080 * 3 / 2)))
    error_check(ret)
    p_ipe_buf_out, ret = sdk.lyn_malloc(model_info.input_size)
    error_check(ret)
    p_dev_infer_buf_out, ret = sdk.lyn_malloc(model_info.output_size)
    error_check(ret)

    print("store in face lib ...")
    with open(feature_path, "wb") as f:
        for image in images:
            image_dec = image_decoder.ImageDecoder()
            image_dec.init(face_image_path + "/" + image, False)
            imgInfo = image_dec.get_image_dec_info()

            decodeImg = sdk.lyn_codec_buf_t()
            decodeImg.size = imgInfo.output.predict_buf_size
            decodeImg.data = p_decode_buf_out
            image_dec.decode_image_to_device(stream, decodeImg.data)

            ipe.set_img_info(
                imgInfo.output.width,
                imgInfo.output.height,
                sdk.lyn_pixel_format_t.LYN_PIX_FMT_NV12,
            )
            ipe.calc_param(stream, decodeImg.data, p_ipe_buf_out)

            p_host_infer_buf_out = sdk.c_malloc(model_info.output_size)
            error_check(
                sdk.lyn_execute_model_async(
                    stream,
                    model_info.model,
                    p_ipe_buf_out,
                    p_dev_infer_buf_out,
                    model_info.batch_size,
                )
            )
            error_check(
                sdk.lyn_memcpy_async(
                    stream,
                    p_host_infer_buf_out,
                    p_dev_infer_buf_out,
                    model_info.output_size,
                    sdk.lyn_memcpy_dir_t.ServerToClient,
                )
            )

            error_check(
                sdk.lyn_stream_add_callback(
                    stream,
                    write_callback,
                    [p_host_infer_buf_out, f, model_info.output_size],
                )
            )
            error_check(sdk.lyn_synchronize_stream(stream))

        print("store over")
        ipe.destory()
        error_check(sdk.lyn_destroy_stream(stream))
        error_check(sdk.lyn_free(p_dev_infer_buf_out))
        error_check(sdk.lyn_free(p_ipe_buf_out))
        error_check(sdk.lyn_free(p_decode_buf_out))
        model_info.unload_model()
        error_check(sdk.lyn_destroy_context(ctx))


if __name__ == "__main__":
    main()
